# -*- coding: utf-8 -*-
"""
@Time    : 2024/8/22 13:33 
@Author  : ZhangShenao 
@File    : gaode_weather_tool.py 
@Desc    : 高德天气查询工具

高德Web服务平台: https://lbs.amap.com/api/webservice/summary
高德开发平台控制台: https://console.amap.com/dev/index
"""
import json
import os
from typing import Any, Type, Dict

import requests
from langchain_core.pydantic_v1 import Field, BaseModel
from langchain_core.tools import BaseTool

BASE_URL = "https://restapi.amap.com/v3"  # 高德API基础URL


class GaodeWeatherArgSchema(BaseModel):
    """高德天气查询工具输入参数"""

    city: str = Field(description="城市名称")


class GaodeWeatherTool(BaseTool):
    """高德天气查询工具"""

    name = "gaode_weather_tool"  # 工具名称
    description = "根据输入的城市,查询天气信息"  # 工具描述
    args_schema: Type[BaseModel] = GaodeWeatherArgSchema  # 输入参数定义
    return_direct = True

    def _run(self, *args: Any, **kwargs: Any) -> str:
        # 获取输入的城市信息
        city = kwargs.get("city")
        if city is None:
            return "输入的城市信息有误"

        # 获取高德api_key
        api_key = os.getenv("GAODE_API_KEY")
        if api_key is None:
            return "请先设置高德api_key"

        try:
            # 根据城市信息,获取城市的区域编码adcode
            city_url = f"{BASE_URL}/config/district?keywords={city}&subdistrict=0&key={api_key}"
            city_resp = self._get_json_response(city_url)
            if city_resp.get("info") != "OK":
                return "获取城市编码失败"
            city_code = city_resp.get("districts")[0].get("adcode")

            # 根据城市编码adcode,查询天气信息
            weather_url = f"{BASE_URL}/weather/weatherInfo?extensions=all&city={city_code}&key={api_key}"
            weather_resp = self._get_json_response(weather_url)
            if weather_resp.get("info") != "OK":
                return "获取天气信息失败"

            # 将数据转换为json格式返回
            return json.dumps(weather_resp)

        except Exception as e:
            return f"查询{city}的天气信息失败,错误信息:{str(e)}"

    @staticmethod
    def _get_json_response(url: str) -> Dict:
        """请求高德API,获取json响应"""

        with requests.session():
            session = requests.session()
            resp = session.request(
                method="GET",
                url=url,
                headers={"Content-Type": "application/json; charset=utf-8"},
            )
            resp.raise_for_status()
            return resp.json()
